#include <iostream>
#include <fstream>
#include <string>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/md5.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <net/if.h>

using namespace std;

#define READ_BUFF 1024

void encryptFile(const string& inputFile, const string& outputFile, const unsigned char key[32], const unsigned char iv[16]) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    const EVP_CIPHER *cipher = EVP_aes_256_cbc();

    EVP_CIPHER_CTX_init(ctx);
    EVP_EncryptInit_ex(ctx, cipher, NULL, key, iv);

    ifstream in(inputFile, ifstream::in | ios::binary);
    if (!in) {
        throw runtime_error("Failed to open input file.");
    }

    ofstream out(outputFile, ios::binary);
    if (!out) {
        throw runtime_error("Failed to create output file.");
    }

    int out_len = 0;
    char inbuf[READ_BUFF];
    char outbuf[READ_BUFF + 520];

    while (!in.eof()) {
        in.read(inbuf, READ_BUFF);
        int read_len = in.gcount();
        if (!EVP_EncryptUpdate(ctx, (unsigned char*)outbuf, &out_len, (unsigned char*)inbuf, read_len)) {
            cerr << "EVP_EncryptUpdate failed" << endl;
            throw runtime_error("Encryption failed.");
        }
        out.write(outbuf, out_len);
    }
    in.close();

    if (!EVP_EncryptFinal_ex(ctx, (unsigned char*)outbuf, &out_len)) {
        cerr << "EVP_EncryptFinal_ex failed" << endl;
        throw runtime_error("Final encryption step failed.");
    }
    out.write(outbuf, out_len);
    out.close();

    EVP_CIPHER_CTX_cleanup(ctx);
    std::cout << "Encryption successful." << endl;
}
std::vector<char> decryptFileToBinaryStream(const std::string& inputFile, const unsigned char key[32], const unsigned char iv[16])
{
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    const EVP_CIPHER *cipher = EVP_aes_256_cbc();

    EVP_CIPHER_CTX_init(ctx);
    EVP_DecryptInit_ex(ctx, cipher, NULL, key, iv);

    std::ifstream in(inputFile, std::ios::in | std::ios::binary);
    if (!in)
    {
        throw std::runtime_error("Failed to open input file.");
    }

    std::vector<char> decryptedData;
    int out_len = 0;
    char inbuf[READ_BUFF];
    char outbuf[READ_BUFF + 520];

    while (!in.eof())
    {
        in.read(inbuf, READ_BUFF);
        int read_len = in.gcount();
        if (!EVP_DecryptUpdate(ctx, reinterpret_cast<unsigned char*>(outbuf), &out_len, reinterpret_cast<const unsigned char*>(inbuf), read_len))
        {
            std::cerr << "EVP_DecryptUpdate failed" << std::endl;
            throw std::runtime_error("Decryption failed.");
        }
        decryptedData.insert(decryptedData.end(), outbuf, outbuf + out_len);
    }
    in.close();

    if (!EVP_DecryptFinal_ex(ctx, reinterpret_cast<unsigned char*>(outbuf), &out_len))
    {
        std::cerr << "EVP_DecryptFinal_ex failed" << std::endl;
        throw std::runtime_error("Final decryption step failed.");
    }
    decryptedData.insert(decryptedData.end(), outbuf, outbuf + out_len);

    EVP_CIPHER_CTX_cleanup(ctx);
    return decryptedData;
}

void decryptFile(const string& inputFile, const string& outputFile, const unsigned char key[32], const unsigned char iv[16]) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    const EVP_CIPHER *cipher = EVP_aes_256_cbc();

    EVP_CIPHER_CTX_init(ctx);
    EVP_DecryptInit_ex(ctx, cipher, NULL, key, iv);

    ifstream in(inputFile, ifstream::in | ios::binary);
    if (!in) {
        throw runtime_error("Failed to open input file.");
    }

    ofstream out(outputFile, ios::binary);
    if (!out) {
        throw runtime_error("Failed to create output file.");
    }

    int out_len = 0;
    char inbuf[READ_BUFF];
    char outbuf[READ_BUFF + 520];

    while (!in.eof()) {
        in.read(inbuf, READ_BUFF);
        int read_len = in.gcount();
        if (!EVP_DecryptUpdate(ctx, (unsigned char*)outbuf, &out_len, (unsigned char*)inbuf, read_len)) {
            cerr << "EVP_DecryptUpdate failed" << endl;
            throw runtime_error("Decryption failed.");
        }
        out.write(outbuf, out_len);
    }
    in.close();

    if (!EVP_DecryptFinal_ex(ctx, (unsigned char*)outbuf, &out_len)) {
        cerr << "EVP_DecryptFinal_ex failed" << endl;
        throw runtime_error("Final decryption step failed.");
    }
    out.write(outbuf, out_len);
    out.close();

    EVP_CIPHER_CTX_cleanup(ctx);
    std::cout << "Decryption successful." << endl;
}

// 随机生成密钥和初始化向量
void generateKeyAndIV(unsigned char key[32], unsigned char iv[16]) {
    RAND_bytes(key, 32);
    RAND_bytes(iv, 16);
    
    // 保存秘钥
    ofstream keyFile("../key.bin", ios::binary);
    keyFile.write((char*)key, 32);

    keyFile.close();
    // 保存初始化向量
    ofstream ivFile("../iv.bin", ios::binary);
    ivFile.write((char*)iv, 16);
    for (int i = 0; i < 16; i++) {
        std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)iv[i];
    }
    ivFile.close();
    std::cout << "Key and IV generated." << endl;
    std::cout << "Key and IV saved to files." << endl;
}

// 随机生成key，iv值为默认
void generateKey(unsigned char key[32]) {
    RAND_bytes(key, 32);
    
    // 保存秘钥
    ofstream keyFile("../key.bin", ios::binary);
    keyFile.write((char*)key, 32);

    keyFile.close();
    std::cout << "Key generated." << endl;
    std::cout << "Key saved to files." << endl;
}

// 读取密钥和初始化向量,去除换行符
void readKeyAndIV(unsigned char key[32], unsigned char iv[16]) {
    ifstream keyFile("../key.bin", ios::binary);
    if (!keyFile) {
        throw runtime_error("Failed to open key file.");
    }
    keyFile.read((char*)key, 32);
    keyFile.close();

    ifstream ivFile("../iv.bin", ios::binary);
    if (!ivFile) {
        throw runtime_error("Failed to open IV file.");
    }
    ivFile.read((char*)iv, 16);
    ivFile.close();
    std::cout << "Key and IV read from files." << endl;
}

// 只读取key值，iv值为默认
void readKey(unsigned char key[32]) {
    ifstream keyFile("../key.bin", ios::binary);
    if (!keyFile) {
        throw runtime_error("Failed to open key file.");
    }
    keyFile.read((char*)key, 32);
    keyFile.close();
    std::cout << "Key and IV read from files." << endl;
}
